from typing import Callable

import numpy as np
import math
from scipy.interpolate import CubicSpline
import dubins

from _dora_utils import closest_vertex
from _hd_map import HDMap
from dora import DoraStatus
from carla import Map
import math

# Global Planner
GLOBAL_PLANNER_FILL_WITH_DUBINS = True
GLOBAL_PLANNER_SMOOTH_WAYPOINTS = False
GLOBAL_PLANNER_NUM_WAYPOINTS_AHEAD = 120
GLOBAL_PLANNER_TARGET_SPEED = 5.0  # m/s
GLOBAL_PLANNER_TURNING_RADIUS = 1.5  # m
GLOBAL_PLANNER_STEP_SIZE = 1.0  # m
GLOBAL_PLANNER_EXTEND_LEN = 3.0  # m

# Debug
DEBUG = False


class Operator:
    """
    Compute global planning waypoints, given a `position`, an `objective waypoints` and an `opendrive` map.
    """

    def __init__(self):
        self.hd_map = None
        self.position = np.array([])
        self.goal_location = np.array([])
        self.waypoints_use = np.array([])
        self.waypoints_full = np.array([])
        self.target_speeds_use = np.array([])
        self.target_speeds_full = np.array([])

    def on_event(
        self,
        dora_event: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            return self.on_input(dora_event, send_output)
        return DoraStatus.CONTINUE

    def on_input(
        self,
        dora_input: dict,
        send_output: Callable[[str, bytes], None],
    ):
        if dora_input["id"] == "position":
            self.position = np.frombuffer(dora_input["data"], np.float32)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "opendrive":
            if self.hd_map is None:
                opendrive = dora_input["data"].decode()
                self.hd_map = HDMap(
                    Map("map", opendrive), resolution=GLOBAL_PLANNER_STEP_SIZE
                )
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "objective_waypoints":
            objective_waypoints = np.frombuffer(dora_input["data"], np.float32).reshape(
                (-1, 3)
            )

            if self.hd_map is None:
                if DEBUG:
                    print("[Global Planner] No opendrive map found")
                return DoraStatus.CONTINUE

            if len(self.position) == 0:
                if DEBUG:
                    print("[Global Planner] No position found")
                return DoraStatus.CONTINUE

            if len(objective_waypoints) == 0:
                if DEBUG:
                    print(f"[Globa Planner] No objective waypoint found")
                return DoraStatus.CONTINUE

            # Reset if objective waypoint changes
            if not np.array_equal(self.goal_location, objective_waypoints[0]):
                self.goal_location = objective_waypoints[0]
                self.waypoints_use = np.array([])
                self.waypoints_full = np.array([])
                self.target_speeds_use = np.array([])
                self.target_speeds_full = np.array([])

            # Used cached waypoints but start at closest point
            if len(self.waypoints_use) != 0:
                (index, _) = closest_vertex(
                    self.waypoints_full[:, :2],
                    np.array([self.position[:2]]),
                )
                index = max(index - 1, 0)  # lookback one point

                self.waypoints_use = self.waypoints_full[
                    index : min(
                        index + GLOBAL_PLANNER_NUM_WAYPOINTS_AHEAD,
                        len(self.waypoints_full),
                    )
                ]
                self.target_speeds_use = self.target_speeds_full[
                    index : min(
                        index + GLOBAL_PLANNER_NUM_WAYPOINTS_AHEAD,
                        len(self.waypoints_full),
                    )
                ]

            elif len(self.waypoints_use) == 0:
                # Compute the waypoints
                [x, y, z, _, _, _, _] = self.position
                waypoints = self.hd_map.compute_waypoints(
                    [
                        x,
                        y,
                        z,
                    ],
                    self.goal_location,
                )

                # Connect the waypoints using dubins path if they are too far away
                if GLOBAL_PLANNER_FILL_WITH_DUBINS and len(waypoints) > 1:
                    # Compute theta for waypoints
                    thetas = []
                    for i in range(len(waypoints) - 1):
                        dx = waypoints[i + 1][0] - waypoints[i][0]
                        dy = waypoints[i + 1][1] - waypoints[i][1]
                        theta = np.arctan2(dy, dx)
                        thetas.append(theta)
                    thetas.append(thetas[-1])

                    refined_waypoints = np.array([waypoints[0]])

                    for i in range(len(waypoints) - 1):
                        current_pt = np.array(
                            [waypoints[i, 0], waypoints[i, 1], thetas[i]]
                        )
                        next_pt = np.array(
                            [waypoints[i + 1, 0], waypoints[i + 1, 1], thetas[i + 1]]
                        )
                        distance = np.linalg.norm(next_pt[:2] - current_pt[:2])
                        if distance > GLOBAL_PLANNER_STEP_SIZE * 3:
                            dubins_path = self.compute_dubins_path(
                                current_pt,
                                next_pt,
                                GLOBAL_PLANNER_FILL_WITH_DUBINS,
                                GLOBAL_PLANNER_STEP_SIZE,
                            )
                            dubins_path = dubins_path[:, :2]
                            dubins_path = np.column_stack(
                                (
                                    dubins_path,
                                    np.ones(len(dubins_path))
                                    * (waypoints[i, 2] + waypoints[i + 1, 2])
                                    / 2,
                                )
                            )
                            refined_waypoints = np.concatenate(
                                (refined_waypoints, dubins_path[:-1]), axis=0
                            )
                            if DEBUG:
                                print(
                                    f"[Global Planner] Found two points with a distance of {distance} m. Filled with dubins curve. {len(dubins_path)} new points added."
                                )
                        else:
                            refined_waypoints = np.concatenate(
                                (refined_waypoints, waypoints[i + 1, :].reshape(1, -1)),
                                axis=0,
                            )

                    waypoints = refined_waypoints

                # Smooth the global points
                if GLOBAL_PLANNER_SMOOTH_WAYPOINTS and len(waypoints) > 1:
                    waypoints = self.filter_consecutive_duplicate(waypoints)
                    dx = np.diff(waypoints[:, 0])
                    dy = np.diff(waypoints[:, 1])
                    dz = np.diff(waypoints[:, 2])

                    distances = np.sqrt(dx**2 + dy**2 + dz**2)
                    s = np.concatenate(([0.0], np.cumsum(distances)))

                    spline_x = CubicSpline(s, waypoints[:, 0], bc_type="not-a-knot")
                    spline_y = CubicSpline(s, waypoints[:, 1], bc_type="not-a-knot")
                    spline_z = CubicSpline(s, waypoints[:, 2], bc_type="not-a-knot")

                    resampled_s = np.arange(
                        0, s[-1] + GLOBAL_PLANNER_STEP_SIZE, GLOBAL_PLANNER_STEP_SIZE
                    )
                    smooth_x = spline_x(resampled_s)
                    smooth_y = spline_y(resampled_s)
                    smooth_z = spline_z(resampled_s)

                    waypoints = np.column_stack((smooth_x, smooth_y, smooth_z))

                # extend head and tail waypoints for stability
                if len(waypoints > 1):
                    # extend head
                    extend_dir = waypoints[0, :2] - waypoints[1, :2]
                    extend_dir = extend_dir / np.linalg.norm(extend_dir)
                    num_new_pts = math.ceil(GLOBAL_PLANNER_EXTEND_LEN / GLOBAL_PLANNER_STEP_SIZE)
                    for _ in range(num_new_pts):
                        new_xy = waypoints[0, :2] + extend_dir * GLOBAL_PLANNER_STEP_SIZE
                        waypoints = np.vstack(
                            [np.array([new_xy[0], new_xy[1], waypoints[0, 2]]), waypoints]
                        )
                    # extend tail
                    extend_dir = waypoints[-1, :2] - waypoints[-2, :2]
                    extend_dir = extend_dir / np.linalg.norm(extend_dir)
                    num_new_pts = math.ceil(GLOBAL_PLANNER_EXTEND_LEN / GLOBAL_PLANNER_STEP_SIZE)
                    for _ in range(num_new_pts):
                        new_xy = waypoints[-1, :2] + extend_dir * GLOBAL_PLANNER_STEP_SIZE
                        waypoints = np.vstack(
                            [waypoints, np.array([new_xy[0], new_xy[1], waypoints[-1, 2]])]
                        )

                self.waypoints_full = waypoints.copy()
                self.target_speeds_full = np.array(
                    [GLOBAL_PLANNER_TARGET_SPEED] * len(waypoints)
                )  # currently not used

                (index, _) = closest_vertex(
                    self.waypoints_full[:, :2],
                    np.array([self.position[:2]]),
                )
                index = max(index - 1, 0)  # lookback one point

                self.waypoints_use = self.waypoints_full[
                    index : min(
                        index + GLOBAL_PLANNER_NUM_WAYPOINTS_AHEAD,
                        len(self.waypoints_full),
                    )
                ]
                self.target_speeds_use = self.target_speeds_full[
                    index : min(
                        index + GLOBAL_PLANNER_NUM_WAYPOINTS_AHEAD,
                        len(self.waypoints_full),
                    )
                ]

            if len(self.waypoints_use) == 0:
                send_output(
                    "gps_waypoints",
                    self.waypoints_use.tobytes(),
                    {},
                )
                return DoraStatus.CONTINUE

            # Send output
            global_waypoints = np.column_stack(
                [
                    self.waypoints_use,  # x, y, z
                    self.target_speeds_use,  # target_speed (not used)
                ]
            ).astype(np.float32)
            global_waypoints = self.filter_consecutive_duplicate(global_waypoints)

            send_output(
                "gps_waypoints",
                global_waypoints.tobytes(),
                {},
            )

            return DoraStatus.CONTINUE

    @staticmethod
    def filter_consecutive_duplicate(array: np.ndarray):
        return np.array(
            [elem for i, elem in enumerate(array) if (elem - array[i - 1]).any() or i == 0]
        )

    @staticmethod
    def compute_dubins_path(
        start_pt: np.ndarray,
        end_pt: np.ndarray,
        turning_radius: float = 1.0,
        step_size: float = 0.1,
    ):
        dubins_path = dubins.shortest_path(start_pt, end_pt, turning_radius)
        waypoints, _ = dubins_path.sample_many(step_size)

        return np.asarray(waypoints)
